import numpy as np
import pandas as pd
import json
import os
# 
current_path = os.path.dirname(os.path.abspath(__file__))
os.chdir(current_path)
#
TIME_SERIES_LENGTH = 128


def main():
    dataset_filename = "6U2U3mSAC_length_128.csv"
    # change the lables file to generate different dataset
    # labels_name = "lables_onlyHorizontalParts.json"
    # output_filename = dataset_filename[0:-4]+"_"+labels_name[11:19]+"_dataset"
    
    labels_name = "labels.json"
    output_filename = dataset_filename[0:-4]+ "_VH_dataset"
    

    print("loading data")
    # 1094 rows
    x = pd.read_csv(dataset_filename)
    # y = pd.read_csv("y_train_EFo1WyE.csv")
    # print(x)

    print("loading labels.json")
    # Load labels, file can be found in challenge description
    with open(labels_name, "r") as stream_json:
        labels = json.load(stream_json)
    print(labels)
    
    m = x.shape[0]
    K = TIME_SERIES_LENGTH  # Can be found through csv

    # Create R and Z
    print("------------------------------")
    print(labels["R"])
    R = x[[f"{var_name}_{i}" for var_name in labels["R"] for i in range(K)]]
    R = R.values.reshape((m, -1, K))
    R = R.astype(np.float32)
    print("R.shape :  ", R.shape)
    # print(R)

    print("------------------------------")
    # print(labels["Z"])
    Z = x[[f"{var_name}_{i}" for var_name in labels["Z"] for i in range(K)]]
    Z = Z.values.reshape((m, -1, K))
    Z = Z.astype(np.float32)
    print("Z.shape :  ", Z.shape)    
    # print(Z)

    print("------------------------------")
    # X = x[[f"{var_name}_{i}" for var_name in labels["X"] for i in range(K)]]
    X = x[[f"{var_name}_{i}" for var_name in labels["X"] for i in range(K)]]
    X = X.values.reshape((m, -1, K))
    X = X.astype(np.float32)
    print("X.shape :  ", X.shape)    
    # print(X)
    # 
    print("------------------------------")


    print("saving structures....")
    np.savez(output_filename, R=R, X=X, Z=Z)
    print(output_filename+'.npz')


if __name__ == "__main__":
    main()
